import base64
from vlmeval.smp import *
from vlmeval.api.base import BaseAPI
from vlmeval.dataset import DATASET_TYPE
from vlmeval.smp.vlm import encode_image_file_to_base64
import time


class bailingMMWrapper(BaseAPI):

    is_api: bool = True

    def __init__(
        self,
        model: str,
        retry: int = 5,
        wait: int = 5,
        key: str = None,
        verbose: bool = True,
        system_prompt: str = None,
        max_tokens: int = 1024,
        proxy: str = None,
        **kwargs,
    ):

        self.model = model
        self.fail_msg = "Failed to obtain answer via bailingMM API."
        if key is None:
            key = os.environ.get("BAILINGMM_API_KEY", None)
        assert key is not None, "Please set the API Key for bailingMM."
        self.key = key
        self.headers = {"Content-Type": "application/json"}
        super().__init__(
            wait=wait,
            retry=retry,
            system_prompt=system_prompt,
            verbose=verbose,
            **kwargs,
        )

    def image_to_base64(self, image_path):
        with open(image_path, "rb") as image_file:
            encoded_string = str(base64.b64encode(image_file.read()), "utf-8")
            return encoded_string

    def prepare_inputs(self, inputs):
        msgs = cp.deepcopy(inputs)
        content = []
        for i, msg in enumerate(msgs):
            if msg["type"] == "text":
                pass
            else:
                try:
                    image_data = self.image_to_base64(msg["value"])
                except Exception as e:
                    if self.verbose:
                        self.logger.error(e)
                    image_data = ""
                msg["value"] = image_data
            content.append(msg)
        return content

    def generate_inner(self, inputs, **kwargs) -> str:
        assert isinstance(inputs, str) or isinstance(inputs, list)
        start = time.time()
        inputs = [inputs] if isinstance(inputs, str) else inputs

        messages = self.prepare_inputs(inputs)

        service_url = "https://bailingchat.alipay.com/api/proxy/eval/antgmm/completions"

        payload = {
            "structInput": messages,
            "sk": self.key,
            "model": self.model,
            "timeout": 180000,
        }
        response = requests.post(service_url, headers=self.headers, json=payload)
        if self.verbose:
            self.logger.info("Time for requesting is:")
            self.logger.info(time.time() - start)
        try:
            assert response.status_code == 200
            output = json.loads(response.text)
            answer = output["preds"]["pred"]
            if self.verbose:
                self.logger.info(f"inputs: {inputs}\nanswer: {answer}")
            return 0, answer, "Succeeded! "
        except Exception as e:
            if self.verbose:
                self.logger.error(e)
                self.logger.error(f"The input messages are {inputs}.")
            return -1, self.fail_msg, ""


class bailingMMAPI(bailingMMWrapper):

    def generate(self, message, dataset=None):
        return super(bailingMMAPI, self).generate(message, dataset=dataset)
